Skip to content

[MoE][Common/PyTorch] Add permutation#936

Merged
phu0ngng merged 35 commits intoNVIDIA:mainfrom
StudyingShao:jiangs/permutation
Aug 22, 2024
Merged

[MoE][Common/PyTorch] Add permutation#936
phu0ngng merged 35 commits intoNVIDIA:mainfrom
StudyingShao:jiangs/permutation

Conversation

@StudyingShao
Copy link
Copy Markdown
Contributor

@StudyingShao StudyingShao commented Jun 17, 2024 •
edited by phu0ngng
Loading

Description

Permutation for fp32/bf16/fp16/fp8 data type. Now PyTorch op only.

Additional descriptions: https://github.com/fanshiqing/moe_grouped_gemm/tree/dev

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refractor

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@StudyingShao StudyingShao changed the title [Draft][Common/PyTorch] Add permutation [Draft][MoE][Common/PyTorch] Add permutation Jun 24, 2024
@StudyingShao StudyingShao force-pushed the jiangs/permutation branch 3 times, most recently from 225a5be to 4b24d86 Compare June 25, 2024 20:04
@StudyingShao @QiZhangNV
Add submodule cutlass v3.5.0
2989e5c 
Co-authored-by: Qi Zhang <qizhang@nvidia.com>
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Copy link
Copy Markdown
Contributor Author

StudyingShao commented Jun 25, 2024

Hi @phu0ngng @cyanguwa , this PR is the Permutation fusion operators needed by MoE.
Please ignore the unit test file tests/pytorch/test_permutation.py, and help to review other changes. Thanks.
I will start to refactor the unit test file in parallel.

cc @QiZhangNV

@StudyingShao StudyingShao changed the title [Draft][MoE][Common/PyTorch] Add permutation [MoE][Common/PyTorch] Add permutation Jun 28, 2024
@phu0ngng phu0ngng requested review from cyanguwa and phu0ngng June 28, 2024 16:02
@phu0ngng phu0ngng added the enhancement New feature or request label Jun 28, 2024
@phu0ngng
Copy link
Copy Markdown
Collaborator

phu0ngng commented Jun 28, 2024

/te-ci pytorch

@phu0ngng
Copy link
Copy Markdown
Collaborator

phu0ngng commented Jul 1, 2024

Hi @StudyingShao, thanks for putting this work into TE.
I have a couple of suggestions after the first glance at your code.

  1. Please sign off all of your commits (DCO failed).
  2. Please rewrite the unit test with pytest and enable skipping if FP8 is unavailable (see

    TransformerEngine/tests/pytorch/test_numerics.py

    Line 495 in 7326af9

    pytest.skip(reason_for_no_fp8)
    ).

phu0ngng
phu0ngng reviewed Jul 1, 2024
View reviewed changes
Comment thread transformer_engine/common/include/transformer_engine/permutation.h Outdated
phu0ngng
phu0ngng requested changes Jul 2, 2024
View reviewed changes
Comment thread transformer_engine/common/permutation/permutation.cu Outdated
phu0ngng
phu0ngng reviewed Jul 3, 2024
View reviewed changes
Comment thread transformer_engine/pytorch/permutation.py Outdated
Comment thread transformer_engine/pytorch/permutation.py Outdated
Comment thread transformer_engine/pytorch/csrc/extensions/pybind.cpp Outdated
Comment thread transformer_engine/pytorch/csrc/extensions/permutation.cu Outdated
Comment thread transformer_engine/pytorch/csrc/extensions/permutation.cu Outdated
phu0ngng
phu0ngng reviewed Jul 3, 2024
View reviewed changes
Comment thread transformer_engine/common/permutation/permutation.cu Outdated
phu0ngng
phu0ngng reviewed Jul 3, 2024
View reviewed changes
Comment thread transformer_engine/common/permutation/permutation.cu Outdated
Comment thread transformer_engine/common/permutation/permutation.cu Outdated
Comment thread transformer_engine/common/permutation/permutation.cu Outdated
yaox12
yaox12 reviewed Jul 4, 2024
View reviewed changes
Comment thread transformer_engine/pytorch/csrc/extensions/permutation.cu Outdated
StudyingShao added 10 commits July 9, 2024 08:48
@StudyingShao
Add permutation functions
0dc63b2 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Building pass
4b4930a 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Add permutation ops
5f881d9 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Everything works fine
1ce7fc1 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Replace get_ptr with getDataPtr
fb483f8 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Rename some functions and kernels
c8f7991 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Refactor to fit the TE style. Part I
6b6eb39 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Remove the dependency on cutlass
7135612 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Refactor to fit the TE style. Part II
1f92e83 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Move permutation.py out of module dir
64986fb 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Clear up the code path
31cfa79 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Minor Changes
b37533e 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Add some comments
c8a9977 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Add some comments
83ddfad 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Revise funcion description
f049343 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
phu0ngng
phu0ngng requested changes Jul 26, 2024
View reviewed changes
Comment thread transformer_engine/pytorch/csrc/extensions/permutation.cu Outdated
Comment thread transformer_engine/common/permutation/permutation.cu Outdated
Comment thread tests/pytorch/test_permutation.py Outdated
@StudyingShao
Take NVTETensor as inputs
2b331d4 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Split unit tests
6d23c98 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@phu0ngng phu0ngng requested review from timmoon10 and removed request for cyanguwa August 2, 2024 23:06
timmoon10
timmoon10 reviewed Aug 5, 2024
View reviewed changes
Comment thread transformer_engine/pytorch/permutation.py
Comment thread transformer_engine/pytorch/permutation.py Outdated
@timmoon10 timmoon10 self-requested a review August 5, 2024 22:02
yaox12
yaox12 reviewed Aug 6, 2024
View reviewed changes
Comment thread transformer_engine/pytorch/permutation.py Outdated
@StudyingShao
Change names
0e01f8e 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Use Float8Tensor for FP8 input
906dbb8 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Reformat
6740781 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Rescale fp8 for permute backward
f5ad7fc 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Move dtype to ctx
5307a01 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
phu0ngng
phu0ngng approved these changes Aug 21, 2024
View reviewed changes
Copy link
Copy Markdown
Collaborator

@phu0ngng phu0ngng left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@phu0ngng
Copy link
Copy Markdown
Collaborator

phu0ngng commented Aug 21, 2024

/te-ci pytorch

timmoon10
timmoon10 approved these changes Aug 21, 2024
View reviewed changes
Copy link
Copy Markdown
Collaborator

@timmoon10 timmoon10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM once the CI passes.

StudyingShao and others added 2 commits August 22, 2024 10:14
@StudyingShao
Format fix
318caf8 
Signed-off-by: Jiang Shao <jiangs@nvidia.com>
@StudyingShao
Merge branch 'main' into jiangs/permutation
5f04315
@StudyingShao
Copy link
Copy Markdown
Contributor Author

StudyingShao commented Aug 22, 2024

/te-ci pytorch

@phu0ngng phu0ngng merged commit a335374 into NVIDIA:main Aug 22, 2024
This was referenced Apr 21, 2026
Open
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@timmoon10 timmoon10 timmoon10 approved these changes

@phu0ngng phu0ngng phu0ngng approved these changes

+1 more reviewer

@yaox12 yaox12 yaox12 left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

enhancement New feature or request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@StudyingShao @phu0ngng @yaox12 @timmoon10